package com.cvai.service.impl;

import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.cvai.aiservice.ChatAssistant;
import com.cvai.aiservice.FileChatAssistant;
import com.cvai.aiservice.ImageChatAssistant;
import com.cvai.domain.dto.ChatDto;
import com.cvai.domain.po.Chat;
import com.cvai.domain.po.File;
import com.cvai.domain.po.Memory;
import com.cvai.mapper.ChatMapper;
import com.cvai.mapper.FileMapper;
import com.cvai.mapper.MemoryMapper;
import com.cvai.service.IChatService;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import com.cvai.service.IFileService;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentParser;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.parser.apache.tika.ApacheTikaDocumentParser;
import dev.langchain4j.data.document.splitter.DocumentByCharacterSplitter;
import dev.langchain4j.data.message.Content;
import dev.langchain4j.data.message.ImageContent;
import dev.langchain4j.data.message.TextContent;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.rag.DefaultRetrievalAugmentor;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.rag.query.transformer.CompressingQueryTransformer;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingStore;

import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.BeanUtils;

import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.scheduler.Schedulers;

import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.nio.file.Files;
import java.time.LocalDateTime;
import java.util.ArrayList;
import java.util.Base64;
import java.util.List;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;

/**
 * <p>
 * 服务实现类
 * </p>
 *
 * @author cv
 * @since 2025-03-09
 */
@Slf4j
@Service
@RequiredArgsConstructor
public class ChatServiceImpl extends ServiceImpl<ChatMapper, Chat> implements IChatService {

    private final RedisTemplate<String, String> redisTemplate;
    private final ChatAssistant chatAssistant;
    private final ImageChatAssistant imageChatAssistant;
    private final MemoryMapper memoryMapper;
    private final FileMapper fileMapper;
    private final EmbeddingModel embeddingModel;
//    private final EmbeddingStore embeddingStore;
    private final ChatLanguageModel chatLanguageModel;
    private final StreamingChatLanguageModel streamingChatLanguageModel;
    private final IFileService fileService;

    @Override
    public Flux<String> chat(ChatDto chatDto) {
        log.info("chatByFlux: chatDto:{}", chatDto);
        if (chatDto == null)
            throw new RuntimeException("chatDto不能为空");
        if (chatDto.getUserId() == null) {
            log.info("userId is null");
            redisTemplate.opsForValue().set(chatDto.getMemoryId().toString(), "1", 1, TimeUnit.DAYS);
        }
        // 创建线程安全的字符串容器
        AtomicReference<StringBuilder> fullContent = new AtomicReference<>(new StringBuilder());
        if (chatDto.getType() == 0) {
            log.info("文本");
            return chatAssistant.chatByFlux(chatDto.getMemoryId(), chatDto.getQuestion())
                    // 1. 处理每个数据块（发送到前端 + 本地收集）
                    .doOnNext(s -> {
                        fullContent.get().append(s);
                        log.info("收到流式块: {}", s);
                    })
                    //2. 格式化为sse 格式
                    // .map(s -> "data: " + s + "\n\n")
                    //3. 流结束时异步保存数据
                    .concatWith(Mono.just("end"))
                    .doOnComplete(() -> {
                        String ans = fullContent.get().toString();
                        log.info("收到完整的流式块: {}", ans);
                        if (!ans.isEmpty()) {
                            Mono.fromRunnable(() -> {
                                        if (chatDto.getUserId() != null) {
                                            Chat chat = new Chat();
                                            BeanUtils.copyProperties(chatDto, chat);
                                            chat.setAnswer(ans);
                                            Memory memory = memoryMapper.selectById(chatDto.getMemoryId());
                                            if (memory != null && memory.getIsEdit() == 0) {
                                                //可以改
                                                String s = null;
                                                if(chatDto.getQuestion().length()>50) {
                                                    s = chatDto.getQuestion().substring(0, 50);
                                                }
                                                memory.setName(s == null ?chatDto.getQuestion():s);
                                                memory.setUserId(chatDto.getUserId());
                                                memoryMapper.updateById(memory);
                                            }
                                            chat.setCreateTime(LocalDateTime.now());
                                            save(chat);
                                        }
                                    }).subscribeOn(Schedulers.boundedElastic()) // 切换到阻塞线程池
                                    .subscribe(
                                            null,
                                            error -> log.error("存储失败", error)
                                    );
                        }
                    });
        } else if (chatDto.getType() == 1) {//传入的是图片
            log.info("图片");
            LambdaQueryWrapper<File> qw = new LambdaQueryWrapper<>();
            qw.in(File::getId, chatDto.getFileId());
            List<File> files = fileMapper.selectList(qw);
            List<Content> contents = new ArrayList<>();
            for (File file : files) {
                if (file.getIsPicture() == 1) {
                    java.io.File image = new java.io.File(file.getFilePath());
                    String name = image.getName();
                    String s = name.substring(name.lastIndexOf(".") + 1);
                    String base64 = fileToBase64(image);
                    contents.add(ImageContent.from(base64, "image/" + s));
                    log.info(s);
                } else {
                    throw new RuntimeException("图片和文件不能一起上传");
                }
            }
            contents.add(TextContent.from(chatDto.getQuestion()));
            log.info("contents:{}",contents);
            return imageChatAssistant.chat(chatDto.getMemoryId(), UserMessage.from(contents))
                    // 1. 处理每个数据块（发送到前端 + 本地收集）
                    .doOnNext(s -> {
                        fullContent.get().append(s);
                        log.info("收到流式块: {}", s);
                    })
                    //2. 格式化为sse 格式
                    // .map(s -> "data: " + s + "\n\n")
                    //3. 流结束时异步保存数据
                    .concatWith(Mono.just("end"))
                    .doOnComplete(() -> {
                        String ans = fullContent.get().toString();
                        log.info("收到完整的流式块: {}", ans);
                        if (!ans.isEmpty()) {
                            Mono.fromRunnable(() -> {
                                        if (chatDto.getUserId() != null) {
                                            Chat chat = new Chat();
                                            BeanUtils.copyProperties(chatDto, chat);
                                            chat.setAnswer(ans);
                                            Memory memory = memoryMapper.selectById(chatDto.getMemoryId());
                                            if (memory != null && memory.getIsEdit() == 0) {
                                                //可以改
                                                String s = null;
                                                if(chatDto.getQuestion().length()>50) {
                                                    s = chatDto.getQuestion().substring(0, 50);
                                                }
                                                memory.setName(s == null ?chatDto.getQuestion():s);
                                                memory.setUserId(chatDto.getUserId());
                                                memoryMapper.updateById(memory);
                                            }
                                            chat.setCreateTime(LocalDateTime.now());
                                            save(chat);
                                            for (File file : files) {
                                                file.setChatId(chat.getId());
                                                fileService.updateById(file);
                                            }
                                        }
                                    }).subscribeOn(Schedulers.boundedElastic()) // 切换到阻塞线程池
                                    .subscribe(
                                            null,
                                            error -> log.error("存储失败", error)
                                    );
                        }
                    });
        } else {//传文件
            log.info("文件");
            LambdaQueryWrapper<File> qw = new LambdaQueryWrapper<>();
            qw.in(File::getId, chatDto.getFileId());
            List<File> files = fileMapper.selectList(qw);
            DocumentParser parser = new ApacheTikaDocumentParser();
            EmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();
            //读文件
            for (File file : files) {
                if (file.getIsPicture() == 1)
                    throw new RuntimeException("图片和文件不能一起上传");
                Document document = FileSystemDocumentLoader.loadDocument(file.getFilePath(), parser);
                DocumentSplitter splitter = new DocumentByCharacterSplitter(100, 50);
                //把文件向量化存储
                splitter.split(document).forEach(textSegment -> embeddingStore.add(embeddingModel.embed(textSegment).content(), textSegment));
            }

            EmbeddingStoreContentRetriever retriever = EmbeddingStoreContentRetriever.builder()
                    .embeddingStore(embeddingStore)
                    .embeddingModel(embeddingModel)
                    .maxResults(10)
                    .minScore(0.5)
                    .build();

            CompressingQueryTransformer queryTransformer = new CompressingQueryTransformer(chatLanguageModel);

            FileChatAssistant fileChatAssistant = AiServices.builder(FileChatAssistant.class)
                    .chatMemoryProvider(id -> MessageWindowChatMemory.withMaxMessages(5))
                    .chatLanguageModel(chatLanguageModel)
                    .streamingChatLanguageModel(streamingChatLanguageModel)
                    .retrievalAugmentor(DefaultRetrievalAugmentor.builder()
                            .contentRetriever(retriever)
                            .queryTransformer(queryTransformer)
                            .build()).build();
            return fileChatAssistant.chat(chatDto.getMemoryId(), chatDto.getQuestion())
                    // 1. 处理每个数据块（发送到前端 + 本地收集）
                    .doOnNext(s -> {
                        fullContent.get().append(s);
                        log.info("收到流式块: {}", s);
                    })
                    //2. 格式化为sse 格式
                    // .map(s -> "data: " + s + "\n\n")
                    //3. 流结束时异步保存数据
                    .concatWith(Mono.just("end"))
                    .doOnComplete(() -> {
                        String ans = fullContent.get().toString();
                        log.info("收到完整的流式块: {}", ans);
                        if (!ans.isEmpty()) {
                            Mono.fromRunnable(() -> {
                                        if (chatDto.getUserId() != null) {
                                            Chat chat = new Chat();
                                            BeanUtils.copyProperties(chatDto, chat);
                                            chat.setAnswer(ans);
                                            Memory memory = memoryMapper.selectById(chatDto.getMemoryId());
                                            if (memory != null && memory.getIsEdit() == 0) {
                                                //可以改
                                                String s = null;
                                                if(chatDto.getQuestion().length()>50) {
                                                    s = chatDto.getQuestion().substring(0, 50);
                                                }
                                                memory.setName(s == null ?chatDto.getQuestion():s);
                                                memory.setUserId(chatDto.getUserId());
                                                memoryMapper.updateById(memory);
                                            }
                                            chat.setCreateTime(LocalDateTime.now());
                                            save(chat);
                                            for (File file : files) {
                                                file.setChatId(chat.getId());
                                                fileService.updateById(file);
                                            }
                                        }
                                    }).subscribeOn(Schedulers.boundedElastic()) // 切换到阻塞线程池
                                    .subscribe(
                                            null,
                                            error -> log.error("存储失败", error)
                                    );
                        }
                    });
        }
    }

    @Override
    public List<Chat> listByMemoryId(String memoryId) {
        if (memoryId == null) {
            throw new RuntimeException("memoryId不能为空");
        }
        QueryWrapper<Chat> qw = new QueryWrapper<>();
        qw.lambda().eq(Chat::getMemoryId, memoryId);
        return list(qw);
    }

    public static String fileToBase64(java.io.File file) {
        // 用于存储文件的字节数据
        byte[] fileContent = null;
        try (InputStream inputStream = new FileInputStream(file)) {
            // 获取文件的长度
            fileContent = new byte[(int) file.length()];
            // 读取文件内容到字节数组
            inputStream.read(fileContent);
        } catch (IOException e) {
            // 若读取文件时发生异常，打印异常信息
            e.printStackTrace();
        }

        if (fileContent != null) {
            // 使用 Base64 编码器对字节数组进行编码
            return Base64.getEncoder().encodeToString(fileContent);
        }
        return null;
    }
}
